Clear TODO left:

1. Write Tests calls.
2. Add Docs.
3. struct to store results for ease of analysis.

Weak SDEPINN Solver

The current implementation of SDEPINN is a Physics-Informed Neural Network (PINN) solver for scalar SDEs, solving Fokker–Planck (Kolmogorov forward) PDEs instead of sampling individual trajectories.

1. Problem

We consider a 1D Itô SDE:

$$ dX_t = f(X_t, t), dt + g(X_t, t), dW_t, $$

with scalar state $(X_t \in \mathbb{R})$, general drift (f) and diffusion (g), and initial condition $(X_{t_0} = x_0)$.

The solver learns the probability density (p(x,t)) satisfying the Fokker–Planck PDE:

$$ \frac{\partial p(x,t)}{\partial t} = -\frac{\partial}{\partial x}\Big[f(x,t) p(x,t)\Big] + \frac{1}{2} \frac{\partial^2}{\partial x^2} \Big[g(x,t)^2 p(x,t)\Big]. $$

2. Domain and boundary conditions

Spatial domain: $(x \in [x_0, x_\text{end}])$ - The user must decide this based on rough range of expected SDE solution.
Initial condition approximated as a narrow PDF:

$$ p(x, t_0) \approx \mathcal{N}(x_0, \sigma_{bc}^2) $$

Boundary conditions:

1. Absorbing:
   $p(x_0, t) = p(x_{end}, t) = 0 $

2. Reflecting (zero flux at chosen spatial domain's boundaries) :
   $J(x,t) = f(x,t)p(x,t) - \frac{1}{2} \frac{\partial}{\partial x}\big[g(x,t)^2 p(x,t)\big] = 0$

Currently (for below results) only reflective BCs were enabled - in practice, one can choose the consistent BC loss terms.

3. Loss formulation

The total loss consists of :

1. PDE residual (enforces Fokker–Planck equation):

$$ \mathcal{L}_\text{PDE} = \sum_{(x,t) \in \text{grid}} \Bigg| \frac{\partial p_\theta(x,t)}{\partial t} + \frac{\partial}{\partial x}[f(x,t)p_\theta(x,t)] - \frac{1}{2} \frac{\partial^2}{\partial x^2}[g(x,t)^2 p_\theta(x,t)] \Bigg|^2 $$

2. Boundary condition loss:

$$ \mathcal{L}_\text{BC} = \sum_{x \in {x_0, x_\text{end}}} |p_\theta(x,t)|^2 + |J_\theta(x,t)|^2 $$

3. Normalization loss (probability mass = 1):

$$ \mathcal{L}_\text{norm} = \sum_{t \in \text{grid}} \Big| \int_{x_0}^{x_\text{end}} p_\theta(x,t), dx - 1 \Big|^2 $$

6. Initial condition loss:

$$ \mathcal{L}_\text{IC} = \sum_{x \in \text{IC points}} |p_\theta(x,t_0) - p_\text{IC}(x)|^2 $$

The total loss minimized during training:

$$ \mathcal{L}_\text{total} = \mathcal{L}_\text{PDE} + \mathcal{L}_\text{IC} + \mathcal{L}_\text{BC} + \lambda*\text{norm} \mathcal{L}_\text{norm} $$

Currently, IC loss is implemented pointwise; normalization enforces integrated probability mass.

4. Potential improvements

• Multidimensional SDEs: currently only scalar (X_t); generalizing requires multi-dimensional PDEs with vector X.
• Boundary condition handling: give user choice with absorbing &/or reflecting BC.
• Batching & autodiff: ensure consistency for large grids.
• Parameter estimation mode: ?

TL;DR

SDEPINN solves scalar SDEs at the density level via PINNs, enforcing PDE, BCs, ICs, and normalization.

Current results :

1. GBM SDE (resulting state follows a LogNormal Distribution)

2. OU Process (resulting state follows a Normal distribution)

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.